package main

import (
    "adkr/pkg/adkr"
    "adkr/pkg/config"
    "adkr/pkg/crypto"
    "bytes"
    "encoding/json"
    "flag"
    "fmt"
    "io"
    "log"
    "os"
    "strings"
    "sync"
    "sync/atomic"
    "time"
)

type bus struct {
    n       int
    ins     []chan adkr.Message
    outs    []chan adkr.Message
    stopCh  chan struct{}
    wg      sync.WaitGroup
    sentCnt []int64
    recvCnt []int64
}

func newBus(n int) *bus {
    b := &bus{
        n:      n,
        ins:    make([]chan adkr.Message, n),
        outs:   make([]chan adkr.Message, n),
        stopCh: make(chan struct{}),
        sentCnt: make([]int64, n),
        recvCnt: make([]int64, n),
    }
    for i := 0; i < n; i++ {
        b.ins[i] = make(chan adkr.Message, 4096)
        b.outs[i] = make(chan adkr.Message, 4096)
    }
    return b
}

func (b *bus) start() {
    for i := 0; i < b.n; i++ {
        b.wg.Add(1)
        idx := i
        go func() {
            defer b.wg.Done()
            for {
                select {
                case <-b.stopCh:
                    return
                case msg := <-b.outs[idx]:
                    // route
                    if msg.To >= 0 && msg.To < b.n {
                        msg.Sender = idx
                        atomic.AddInt64(&b.sentCnt[idx], 1)
                        atomic.AddInt64(&b.recvCnt[msg.To], 1)
                        b.ins[msg.To] <- msg
                    } else {
                        // broadcast
                        for j := 0; j < b.n; j++ {
                            // deliver to all (including self) for simplicity
                            copyMsg := msg
                            copyMsg.Sender = idx
                            atomic.AddInt64(&b.sentCnt[idx], 1)
                            atomic.AddInt64(&b.recvCnt[j], 1)
                            b.ins[j] <- copyMsg
                        }
                    }
                }
            }
        }()
    }
}

func (b *bus) stop() {
    close(b.stopCh)
    b.wg.Wait()
    for i := 0; i < b.n; i++ {
        close(b.ins[i])
        close(b.outs[i])
    }
}

// snapshot returns per-node sent/recv counters
func (b *bus) snapshot() (sent []int64, recv []int64) {
    sent = make([]int64, b.n)
    recv = make([]int64, b.n)
    for i := 0; i < b.n; i++ {
        sent[i] = atomic.LoadInt64(&b.sentCnt[i])
        recv[i] = atomic.LoadInt64(&b.recvCnt[i])
    }
    return
}

// metricsCollector intercepts logs and aggregates coin metrics and timeouts
type metricsCollector struct {
    w io.Writer
    mu sync.Mutex
    coin map[int]map[string]interface{}
    timeouts int
}

func newMetricsCollector(w io.Writer) *metricsCollector {
    return &metricsCollector{ w: w, coin: make(map[int]map[string]interface{}) }
}

func (mc *metricsCollector) Write(p []byte) (int, error) {
    // mirror to underlying writer
    if mc.w != nil { _, _ = mc.w.Write(p) }
    // parse line
    line := string(bytes.TrimSpace(p))
    if strings.Contains(line, "ThresholdCoin-METRICS ") {
        idx := strings.Index(line, "ThresholdCoin-METRICS ")
        if idx >= 0 {
            js := strings.TrimSpace(line[idx+len("ThresholdCoin-METRICS "):])
            var m map[string]interface{}
            if json.Unmarshal([]byte(js), &m) == nil {
                // expect pid present
                if v, ok := m["pid"].(float64); ok {
                    pid := int(v)
                    mc.mu.Lock()
                    mc.coin[pid] = m
                    mc.mu.Unlock()
                }
            }
        }
    }
    if strings.Contains(strings.ToLower(line), "timed out") {
        mc.mu.Lock(); mc.timeouts++; mc.mu.Unlock()
    }
    return len(p), nil
}

func main() {
    var n, f, k int
    flag.IntVar(&n, "n", 4, "number of nodes")
    flag.IntVar(&f, "f", 1, "fault tolerance")
    flag.IntVar(&k, "k", 1, "rounds")
    fast := flag.Bool("fast", true, "fast-bench mode without full crypto")
    strongPred := flag.Bool("mvba-strong-predicate", false, "enable strong MVBA predicate (ACSS-verified dealers only)")
    abaEvidence := flag.Bool("aba-evidence", false, "attach supporter evidence in ABA messages")
    abaTick := flag.Int("aba-tick-ms", 10, "ABA polling tick in ms")
    abaRoundTO := flag.Int("aba-round-to-ms", 3000, "ABA per-round timeout in ms")
    krSPBC := flag.Bool("kr-spbc", false, "use SPBC fast-path in Key Refresh (default MVBA)")
    rbcMode := flag.String("rbc-mode", "", "RBC variant: '', 'dual', or 'dual_cp'")
    flag.Parse()

    // install metrics collector as log sink
    mc := newMetricsCollector(os.Stderr)
    log.SetOutput(mc)

    log.Printf("Starting ADKR bench: n=%d f=%d k=%d fast=%v strongPred=%v abaEvidence=%v", n, f, k, *fast, *strongPred, *abaEvidence)
    adkr.SetFastBench(*fast)
    adkr.SetStrongMVBAEnabled(*strongPred)
    adkr.SetABAEvidence(*abaEvidence)
    adkr.SetABATiming(time.Duration(*abaTick)*time.Millisecond, time.Duration(*abaRoundTO)*time.Millisecond)
    adkr.SetKeyRefreshSPBCFast(*krSPBC)
    adkr.SetRBCMode(*rbcMode)
    // fast 模式下默认跳过阈值 coin 聚合验证，便于实验闭环
    adkr.SetCoinSkipVerify(*fast)

    // Prepare crypto materials
    g1 := crypto.GetG1()
    g2 := crypto.GetG2()
    thpkPtr, thpks, sks := crypto.GenerateThresholdKeys(n, f)
    thpk := *thpkPtr
    // single esk for simplicity
    esk := crypto.ZRFromInt(123)
    epks := make([]crypto.G1, n)
    for i := 0; i < n; i++ {
        epks[i] = crypto.MulG1(&g1, &esk)
    }

    // Generate Paillier keys for all nodes and write config files (demo); then load and register
    pubMap := make(map[int]*crypto.PaillierPublicKey)
    privMap := make(map[int]*crypto.PaillierPrivateKey)
    for i := 0; i < n; i++ {
        pk, sk, _ := crypto.GeneratePaillierKeyPair(1024)
        pubMap[i] = pk
        privMap[i] = sk
    }
    _ = config.SavePaillierPublicConfig("paillier_pub.json", pubMap)
    _ = config.SavePaillierPrivateConfig("paillier_priv.json", privMap)
    loadedPub, _ := config.LoadPaillierPublicConfig("paillier_pub.json")
    loadedPriv, _ := config.LoadPaillierPrivateConfig("paillier_priv.json", loadedPub)
    for i := 0; i < n; i++ {
        adkr.RegisterPaillierKeys(i, loadedPub[i], loadedPriv[i])
    }

    // Build bus
    b := newBus(n)
    b.start()
    defer b.stop()

    // Launch N protocol instances
    var wg sync.WaitGroup
    outputs := make([]chan adkr.ADKRHighOutput, n)
    for i := 0; i < n; i++ {
        outputs[i] = make(chan adkr.ADKRHighOutput, 1)
    }

    sid := fmt.Sprintf("adkr-bench-%d", time.Now().Unix())

    start := time.Now()
    for pid := 0; pid < n; pid++ {
        wg.Add(1)
        pid := pid
        go func() {
            defer wg.Done()
            // Construct config and protocol directly to capture statistics
            cfg := adkr.ADKRHighConfig{
                SID:   sid,
                PID:   pid,
                N:     n,
                F:     f,
                K:     k,
                G1:    &g1,
                G2:    &g2,
                EPKs:  epks,
                ESK:   esk,
                THPK:  thpk,
                THPKs: thpks,
                THSK:  sks[pid],
            }
            outCh := make(chan adkr.ADKRHighOutput, 1)
            proto := adkr.NewADKRHighProtocol(cfg, b.ins[pid], b.outs[pid], outCh, log.Default())

            if err := proto.Run(); err != nil {
                log.Printf("PID %d: protocol failed: %v", pid, err)
                return
            }
            select {
            case out := <-outCh:
                outputs[pid] <- out
            case <-time.After(30 * time.Second):
                log.Printf("PID %d: no output within timeout", pid)
            }
        }()
    }

    // Wait for all
    wg.Wait()
    dur := time.Since(start)

    // Print metrics JSON summary
    sentCnt, recvCnt := b.snapshot()
    summary := map[string]interface{}{
        "duration_ms": dur.Milliseconds(),
        "n": n, "f": f, "k": k,
        "timeouts": mc.timeouts,
        "messages_sent": sentCnt,
        "messages_recv": recvCnt,
        "coin_metrics": mc.coin,
    }
    jb, _ := json.Marshal(summary)
    fmt.Println(string(jb))
}


